﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.Text;
using CaptchaLearning.Infrastructure.Model;

namespace CaptchaLearning.Infrastructure.Algorithm
{
    public class BaumWelchLearning: IAlgorithm
    {
        public double Execute(Character character)
        {
            int iterations = Global.CurrentProject.Iterations;

            if (iterations == 0)
                throw new ArgumentException("Iterations cannot be zero.");

            // Baum-Welch algorithm.

            int currentIteration = 1;

            ObservableCollection<double> pi = character.Pi;
            List<List<double>> A = character.TransitionMatrix.Probability;
            List<List<double>> B = character.EmissionMatrix.Probability;

            // Initialization
            double[, ,] epsilon;
            double[,] gamma;

            int K = character.Observations.Count; // K is number of obs, K+2 is number of times
            int T = K;
            int States = character.States.Count;
            epsilon = new double[T, States, States];
            gamma = new double[T, States];
            
            Evaluate eval = new Evaluate();
            do // Until convergence or max iterations is reached
            {
                // 1st step - Calculating the forward probability and the
                //            backward probability for each HMM state.
                double[,] alpha = eval.Forward(character, character.Observations); // alpha = States * T
                double[,] beta = eval.Backward(character, character.Observations); // beta = States * T
                
                // 2nd step - Determining the frequency of the transition-emission pair values
                //            and dividing it by the probability of the entire string.

                // Calculate epsilon values for next computations
                for (int t = 0; t < K-1; t++) // from t=1 to t=K-1
                {
                    for (int i = 0; i < States; i++)
                        for (int j = 0; j < States; j++)
                        {
                            double P = 0; // P(O|Lamda) = alpha[States - 1][K] = sum of the probability in the final column of observations
                            for (int p = 0; p < States; p++)
                                P += alpha[p, K-1];
                            if (P == 0.0) P = 0.25;
                            epsilon[t, i, j] = alpha[i, t] * A[i][j] * beta[j, t + 1] * B[j][t + 1] / P;
                        }
                }
                epsilon[K - 1, States - 1, States - 1] = 1.0;

                // Calculate gamma values for next computations
                for (int t = 0; t < K; t++) // t=1 to K
                {
                    for (int i = 0; i < States; i++)
                    {
                        gamma[t, i] = 0;
                        for (int j = 0; j < States; j++)
                            gamma[t, i] += epsilon[t, i, j];
                    }
                }

                // 3. Continue with parameter re-estimation
                currentIteration++;

                // 3.1 Re-estimation of initial state probabilities 
                for (int i = 0; i < States; i++)
                {
                    pi[i] = gamma[0, i];
                }

                // 3.2 Re-estimation of transition probabilities 

                for (int i = 0; i < States; i++)
                {
                    double fromI = 0; // count number of transitions from state i
                    for (int t = 0; t < K-1; t++)  // from t=1 to t=K-1
                    {
                        fromI += gamma[t, i];
                    }
                    for (int j = 0; j < States; j++) // from j=1 to j=N
                    {
                        if (fromI != 0.0)
                        {
                            double fromI2J = 0; // count number of transitions from state i to j
                            for (int t = 0; t < K - 1; t++) // from t=1 to t=K-1
                            {
                                fromI2J += epsilon[t, i, j];
                            }
                            // calculate probability of a transition from i to j
                            A[i][j] = fromI2J / fromI;
                        }
                        else
                        {
                            A[i][j] = 1.0 / States;
                        }
                    }
                }

                // 3.3 Re-estimation of emission probabilities
                for (int i = 0; i < States; i++) // from j=1 to j=N
                {
                    double atI = 0; // count number of times being at state i
                    for (int t = 0; t < K-1; t++)
                    {
                        atI += gamma[t, i];
                    }
                    for (int j = 0; j < K; j++) // to from the first observation to the last
                    {
                        if (atI != 0.0)
                        {
                            double atISeeJ = 0; // count number of times when state i meets observation j 
                            for (int t = 0; t < K; t++) // from t=1 to t=K
                            {
                                if (character.Observations[t].Equals(character.Observations[j]))
                                    atISeeJ += gamma[t, i]; //epsilon[t, i, j];
                            }
                            // calculate probability of a transition from i to j
                            B[i][j] = atISeeJ / atI;
                        }
                        else
                        {
                            B[i][j] = 1.0 / K;
                        }
                    }
                }

                // Re-assign the new value into the character object
                character.Pi = pi;
                character.TransitionMatrix.Probability = A;
                character.EmissionMatrix.Probability = B;

            } while (currentIteration <= iterations);

            return 0.0d;
        }
    }
}
